///////////////////////////////////////////////////////////////////////////////////////////////////
/// \file formic/fqmc/s_squared.cpp
///
/// \brief   implementation file for the monte carlo spin squared expectation value
///
///////////////////////////////////////////////////////////////////////////////////////////////////

#include<vector>
#include<string>
#include<numeric>
#include<cassert>
#include<algorithm>

#include<boost/format.hpp>
#include<boost/shared_ptr.hpp>

#include<formic/fqmc/s_squared.h>
#include<formic/wfn/configuration.h>
#include<formic/fqmc/vmc_walker.h>
#include<formic/fqmc/sampler.h>
#include<formic/wfn/wfn.h>
#include<formic/input/base.h>
#include<formic/mpi/interface.h>
#include<formic/numeric/numeric.h>
//#include<formic/output/output.h>
#include<formic/exception.h>
//#include<formic/lapack/interface.h>
#include<formic/random/random.h>

namespace formic { namespace fqmc {

  ///////////////////////////////////////////////////////////////////////////////////////////////////
  /// \brief  a class to accumulate the S^2 expectation value during a markov chain evolution
  ///
  ///////////////////////////////////////////////////////////////////////////////////////////////////
  template <class S> class SSquaredAcc : public formic::Accumulator<S> {

    // data members
    private:

      /// \brief  number of atomic orbitals
      int _nao;

//      /// \brief  pointer to the guiding function
//      boost::shared_ptr<formic::CoeffObj<S> > _gf;

      /// \brief  reference to the walker's configuration
      const formic::Configuration & _config;

      /// \brief  the spin squared expectation value
      S _ssquared;

    // private member functions
    private:

    public:

      ///////////////////////////////////////////////////////////////////////////////////////////////////
      /// \brief  constructs the accumulator around a specific coefficient object
      ///
      ///////////////////////////////////////////////////////////////////////////////////////////////////
      SSquaredAcc(const formic::InputBase & userinp,
                  const boost::shared_ptr<formic::CoeffObj<S> > & gf,
                  const boost::shared_ptr<formic::CoeffObj<S> > & vf,
                  const formic::Configuration & config)
        : formic::Accumulator<S>(gf, vf),
          _nao(userinp.get<int>("nao")),
          _config(config),
          _ssquared(formic::zero(S()))
      {}

      /// \brief  return the accumulator's name
      std::string name() const { return "S^2 accumulator"; }

      ///////////////////////////////////////////////////////////////////////////////////////////////////
      /// \brief Accumulate the local contribution to S^2
      ///
      /// \param[in]      userinp  the user's input options
      /// \param[in]      w        the weight to apply to this configuration's contribution
      ///
      ///////////////////////////////////////////////////////////////////////////////////////////////////
      void child_accumulate(const formic::InputBase & userinp, const double w) {

        // get an object to store configuration differences
        formic::ConfigDiff cdiff(userinp);
        cdiff.newly_full().assign(2,0);
        cdiff.newly_empty().assign(2,0);

        // add Sz^2 contribution to the total
        {
          const double sz = 0.5 * _config.a_occ_to_site().size() - 0.5 * _config.b_occ_to_site().size();
          _ssquared += w * this->_vgs * sz * sz * formic::unity(S());
        }

        // get S+S- and S-S+ contributions
        for (int b = 0; b < 2; b++) {

          const int c = 1 - b;

          // get references to occupied and unoccupied orbital lists
          const std::vector<int> & occ_to_site   = ( b == 0 ? _config.a_occ_to_site()     : _config.b_occ_to_site()   );
          const std::vector<int> & unocc_to_site = ( b == 0 ? _config.a_unocc_to_site()   : _config.b_unocc_to_site() );

          // loop over occupieds
          for (std::vector<int>::const_iterator oon = occ_to_site.begin(); oon != occ_to_site.end(); oon++) {

            // skip doubly occupied sites
            if ( _config[2*(*oon)+c] == 1 )
              continue;

            // diagonal term
            _ssquared += 0.5 * w * this->_vgs * formic::unity(S());

            // off-diagonal terms (loop over unoccupieds)
            for (std::vector<int>::const_iterator uon = unocc_to_site.begin(); uon != unocc_to_site.end(); uon++) {

              // skip doubly empty sites
              if ( _config[2*(*uon)+c] == 0 )
                continue;

              // set the configuration difference
              cdiff.newly_full().at(0) = 2 * (*uon) + b;
              cdiff.newly_empty().at(0) = 2 * (*oon) + b;
              cdiff.newly_full().at(1) = 2 * (*oon) + c;
              cdiff.newly_empty().at(1) = 2 * (*uon) + c;

              // get contribution
              _ssquared -= 0.5 * w * this->_vgs * this->_base_vf->get_ratio(cdiff);

            }

          }

        }

      }

      ///////////////////////////////////////////////////////////////////////////////////////////////////
      /// \brief finalizes the accumulation (this may include for example communicating the results)
      ///
      /// \param[in]      userinp  the user's input options
      ///
      ///////////////////////////////////////////////////////////////////////////////////////////////////
      void child_finalize(const formic::InputBase & userinp) {

        const S x = _ssquared / this->_total_weight / this->_vgsa;
        formic::mpi::allreduce(&x, &_ssquared, 1, MPI::SUM);

      }

      ///////////////////////////////////////////////////////////////////////////////////////////////////
      /// \brief  resets the accumulator
      ///
      ///////////////////////////////////////////////////////////////////////////////////////////////////
      void child_reset(const formic::Configuration & config) {
        _ssquared = formic::zero(S());
      }

      ///////////////////////////////////////////////////////////////////////////////////////////////////
      /// \brief  return the value of S^2
      ///
      ///////////////////////////////////////////////////////////////////////////////////////////////////
      S s_squared() const {

        return _ssquared;

      }

  };

} // end namespace fqmc

} // end namespace formic

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief  compute the expectation value of S^2
///
/// \param[in]      userinp        the user's input options
/// \param[in]      sampler        object controlling the sampling
/// \param[in]      vf             a shared pointer to the value function
///
/// \return  the expectation value of S^2
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template<class S> S formic::fqmc::s_squared(const formic::InputBase & userinp,
                                            boost::shared_ptr<formic::fqmc::Sampler<S> > & sampler,
                                            const boost::shared_ptr<formic::CoeffObj<S> > & vf) {

  // create an accumulator for S^2
  boost::shared_ptr<formic::Accumulator<S> > sacc( new formic::fqmc::SSquaredAcc<S>(userinp, sampler->gf(), vf, sampler->config()) );

  // sample
  sampler->simple_sample() << sacc << vf;

  // get a derived class pointer to the S^2 accumulator
  formic::fqmc::SSquaredAcc<S> * sacc_ptr = dynamic_cast<formic::fqmc::SSquaredAcc<S> *>(sacc.get());

  // get the value to return
  const S retval = sacc_ptr->s_squared();
  return retval;

}

template double formic::fqmc::s_squared(const formic::InputBase & userinp,
                                        boost::shared_ptr<formic::fqmc::Sampler<double> > & sampler,
                                        const boost::shared_ptr<formic::CoeffObj<double> > & vf);
template std::complex<double> formic::fqmc::s_squared(const formic::InputBase & userinp,
                                                      boost::shared_ptr<formic::fqmc::Sampler<std::complex<double> > > & sampler,
                                                      const boost::shared_ptr<formic::CoeffObj<std::complex<double> > > & vf);
